# -*- coding: utf-8 -*-
import akshare as ak
import pandas as pd
import numpy as np
import datetime
import os
import argparse
import time
import random
from functools import lru_cache


# 导入安全请求模块
from akshare_request import (
    get_stock_zh_a_hist_safe,
    get_stock_zh_a_spot_em_safe,
    get_tool_trade_date_hist_sina_safe
)
# 配置参数, 特定时间涨势比较好的股票
ANALYSIS_CONFIG = {
    'start_date': '20250407',  # 开始日期，格式：YYYYMMDD
    'end_date': '20250415',    # 结束日期，格式：YYYYMMDD
    'top_n': 100,              # 获取涨幅排名前N的股票
    'min_market_value': 1000000000,  # 最小市值（10亿）
    'max_market_value': 500000000000,  # 最大市值（5000亿）
    'include_st': False,       # 是否包含ST股票
    'include_new': False,      # 是否包含次新股（上市不足1年）
    'output_folder': 'data'    # 输出文件夹
}

def format_date(date_str):
    """格式化日期字符串为YYYYMMDD格式"""
    # 如果输入的是YYYY-MM-DD格式，转换为YYYYMMDD
    if '-' in date_str:
        return date_str.replace('-', '')
    return date_str

def get_stock_list():
    """获取所有A股股票列表"""
    try:
        # 获取A股股票基本信息
        stock_info = get_stock_zh_a_spot_em_safe()
        # 筛选股票代码（沪深A股）
        stock_list = stock_info[stock_info["代码"].str.startswith(("00", "60", "30", "68"))]
        
        # 筛选ST股票
        if not ANALYSIS_CONFIG['include_st']:
            stock_list = stock_list[~stock_list["名称"].str.contains("ST")]
        
        # 筛选市值范围
        stock_list["流通市值"] = pd.to_numeric(stock_list["流通市值"], errors="coerce")
        stock_list = stock_list[
            (stock_list["流通市值"] >= ANALYSIS_CONFIG['min_market_value']) & 
            (stock_list["流通市值"] <= ANALYSIS_CONFIG['max_market_value'])
        ]
        
        return stock_list
    except Exception as e:
        print(f"获取股票列表失败: {e}")
        return pd.DataFrame()

@lru_cache(maxsize=300)
def get_stock_history(stock_code, start_date, end_date, max_retries=3):
    """获取指定股票在给定时间范围内的历史数据，带重试机制"""
    # 格式化日期
    start_date = format_date(start_date)
    end_date = format_date(end_date)
    
    # 重试机制
    for attempt in range(max_retries):
        try:
            # 获取历史数据
            df = get_stock_zh_a_hist_safe(symbol=stock_code, period="daily", 
                                   start_date=start_date, end_date=end_date, 
                                   adjust="qfq")
            
            # 如果不包含次新股且数据不足一年，则返回空DataFrame
            if not ANALYSIS_CONFIG['include_new'] and len(df) < 240:  # 约一年的交易日
                return pd.DataFrame()
                
            return df
        except Exception as e:
            if attempt < max_retries - 1:
                # 随机等待时间，避免同时发起大量请求
                wait_time = random.uniform(0.5, 2.0) * (attempt + 1)
                print(f"{stock_code}: 获取数据失败，第{attempt+1}次重试，等待{wait_time:.1f}秒")
                time.sleep(wait_time)
            else:
                print(f"{stock_code}: 获取历史数据失败 {e}")
                return pd.DataFrame()

def calculate_growth_rate(stock_code, stock_name, start_date, end_date):
    """计算指定时间段内的涨幅"""
    try:
        df = get_stock_history(stock_code, start_date, end_date)
        if df.empty:
            return None
        
        # 获取开始日期和结束日期的收盘价
        start_price = df.iloc[0]["收盘"]
        end_price = df.iloc[-1]["收盘"]
        
        # 计算涨幅
        growth_rate = (end_price - start_price) / start_price * 100
        
        # 计算最大回撤
        max_drawdown = 0
        peak = df.iloc[0]["收盘"]
        for i in range(1, len(df)):
            if df.iloc[i]["收盘"] > peak:
                peak = df.iloc[i]["收盘"]
            else:
                drawdown = (peak - df.iloc[i]["收盘"]) / peak * 100
                if drawdown > max_drawdown:
                    max_drawdown = drawdown
        
        # 计算平均成交量
        avg_volume = df["成交量"].mean()
        
        # 计算波动率（收盘价的标准差/平均值）
        volatility = df["收盘"].std() / df["收盘"].mean() * 100
        
        return {
            "代码": stock_code,
            "名称": stock_name,
            "涨幅(%)": round(growth_rate, 2),
            "最大回撤(%)": round(max_drawdown, 2),
            "平均成交量": int(avg_volume),
            "波动率(%)": round(volatility, 2),
            "开始价格": round(start_price, 2),
            "结束价格": round(end_price, 2),
            "交易天数": len(df)
        }
    except Exception as e:
        print(f"{stock_code}: 计算涨幅失败 {e}")
        return None

def analyze_stocks(start_date, end_date, top_n, batch_size=50):
    """分析指定时间段内涨幅最大的股票，使用批处理减少并发请求"""
    print(f"\n===== 分析 {start_date} 至 {end_date} 区间涨势 =====\n")
    
    # 获取股票列表
    stock_list = get_stock_list()
    if stock_list.empty:
        print("获取股票列表失败，程序结束")
        return pd.DataFrame()
    
    print(f"共获取 {len(stock_list)} 只符合条件的股票")
    
    # 计算每只股票的涨幅，使用批处理
    results = []
    total = len(stock_list)
    
    # 将股票列表分成批次处理
    for batch_start in range(0, total, batch_size):
        batch_end = min(batch_start + batch_size, total)
        batch = stock_list.iloc[batch_start:batch_end]
        
        print(f"\n处理批次: {batch_start+1}-{batch_end}/{total}")
        
        batch_results = []
        for _, row in batch.iterrows():
            stock_code = row["代码"]
            stock_name = row["名称"]
            
            result = calculate_growth_rate(stock_code, stock_name, start_date, end_date)
            if result:
                batch_results.append(result)
        
        # 添加到总结果
        results.extend(batch_results)
        
        # 批次间暂停，避免请求过于频繁
        if batch_end < total:
            pause_time = random.uniform(1.0, 3.0)
            print(f"批次处理完成，暂停 {pause_time:.1f} 秒...")
            time.sleep(pause_time)
    
    # 转换为DataFrame并按涨幅排序
    if not results:
        print("没有符合条件的股票，程序结束")
        return pd.DataFrame()
    
    df_results = pd.DataFrame(results)
    df_results = df_results.sort_values(by="涨幅(%)", ascending=False)
    
    # 获取前N名
    top_stocks = df_results.head(top_n)
    
    return top_stocks

def save_results(df, start_date, end_date):
    """保存结果到文件"""
    if df.empty:
        return
    
    # 创建输出文件夹
    folder_path = os.path.join(os.path.dirname(__file__), ANALYSIS_CONFIG['output_folder'])
    os.makedirs(folder_path, exist_ok=True)
    
    # 生成文件名
    file_name = f"涨幅排名_{start_date}_{end_date}.txt"
    file_path = os.path.join(folder_path, file_name)
    
    # 保存到文件
    df.to_csv(file_path, sep="\t", index=False, encoding="utf-8")
    print(f"\n数据已保存至 {file_path}")

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="分析指定时间区间内涨势最好的股票")
    parser.add_argument("-s", "--start", help="开始日期 (YYYYMMDD 或 YYYY-MM-DD)")
    parser.add_argument("-e", "--end", help="结束日期 (YYYYMMDD 或 YYYY-MM-DD)")
    parser.add_argument("-n", "--top", type=int, help="获取涨幅排名前N的股票")
    parser.add_argument("--min-value", type=float, help="最小市值（亿元）")
    parser.add_argument("--max-value", type=float, help="最大市值（亿元）")
    parser.add_argument("--include-st", action="store_true", help="是否包含ST股票")
    parser.add_argument("--include-new", action="store_true", help="是否包含次新股")
    parser.add_argument("-o", "--output", help="输出文件夹")
    
    args = parser.parse_args()
    
    # 更新配置
    if args.start:
        ANALYSIS_CONFIG['start_date'] = format_date(args.start)
    if args.end:
        ANALYSIS_CONFIG['end_date'] = format_date(args.end)
    if args.top:
        ANALYSIS_CONFIG['top_n'] = args.top
    if args.min_value:
        ANALYSIS_CONFIG['min_market_value'] = args.min_value * 100000000  # 亿元转换为元
    if args.max_value:
        ANALYSIS_CONFIG['max_market_value'] = args.max_value * 100000000  # 亿元转换为元
    if args.include_st:
        ANALYSIS_CONFIG['include_st'] = True
    if args.include_new:
        ANALYSIS_CONFIG['include_new'] = True
    if args.output:
        ANALYSIS_CONFIG['output_folder'] = args.output

def main():
    """主函数"""
    # 解析命令行参数
    parse_arguments()
    
    start_time = time.time()
    print(f"开始分析: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"分析区间: {ANALYSIS_CONFIG['start_date']} 至 {ANALYSIS_CONFIG['end_date']}")
    print(f"筛选条件: 市值范围 {ANALYSIS_CONFIG['min_market_value']/100000000:.1f}亿 - {ANALYSIS_CONFIG['max_market_value']/100000000:.1f}亿")
    print(f"包含ST股票: {'是' if ANALYSIS_CONFIG['include_st'] else '否'}")
    print(f"包含次新股: {'是' if ANALYSIS_CONFIG['include_new'] else '否'}")
    
    try:
        # 分析股票
        top_stocks = analyze_stocks(
            ANALYSIS_CONFIG['start_date'],
            ANALYSIS_CONFIG['end_date'],
            ANALYSIS_CONFIG['top_n']
        )
        
        # 显示结果
        if not top_stocks.empty:
            print(f"\n===== 涨幅排名前 {ANALYSIS_CONFIG['top_n']} 的股票 =====\n")
            print(top_stocks.to_string(index=False))
            
            # 保存结果
            save_results(top_stocks, ANALYSIS_CONFIG['start_date'], ANALYSIS_CONFIG['end_date'])
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    except Exception as e:
        print(f"\n程序执行出错: {e}")
    finally:
        end_time = time.time()
        elapsed_time = end_time - start_time
        print(f"\n分析完成，耗时: {elapsed_time:.2f} 秒 ({elapsed_time/60:.2f} 分钟)")
        print(f"结束时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")


if __name__ == "__main__":
    main()